from random import randint
from time import sleep
import csv
# agent只需要根据qtable选择往哪儿走


class Agent:

    def __init__(self, name, qtable, reward, env, pos, alpha, gamma) -> None:

        self.qtable = qtable
        self.reward = reward

        self.alpha = alpha
        self.gamma = gamma

        self.border = -0.9

        self.steps = 0
        self.episode = 0

        self.size = len(self.reward)
        self.pos = pos
        self.pos_backup = pos
        self.end = (self.size-1, self.size-1)

        self.env = env

        self.name = name

        self.log = open(self.name + '.log', 'w')

    def isEnd(self) -> None:
        return self.pos == self.end

    def reset(self) -> None:
        self.pos = self.pos_backup
        self.steps = 0
        self.env.reset()

    def writeQtable(self, filename) -> None:
        with open(filename, 'w') as f:
            f_csv = csv.writer(f)
            for i in range(self.size):
                for j in range(self.size):
                    f_csv.writerow(self.qtable[i][j])

    def getDirection(self) -> int:
        # 找到当前agent的位置对应的qtable位置
        x, y = self.pos
        directs = self.qtable[x][y]

        # 找到概率最大的，如果最大的不为一个，随机选一个
        qmax = max(directs)
        qmax_index = []
        for i in range(len(directs)):
            if directs[i] == qmax:
                qmax_index.append(i)

        # print('qmax_index:', qmax_index)

        # direct是最终得到的方向，0,1,2,3分别代表上下左右
        direct = qmax_index[randint(0, len(qmax_index)-1)]
        return direct

    def action(self) -> None:

        # 无论碰到墙壁或者边界，都算一次移动
        self.steps += 1
        # 确定一个方向
        direct = self.getDirection()
        # print('name:', self.name, 'direct:', direct)

        x, y = self.pos
        # 在env中判断是否跨越边界
        if self.env.crossBorder(self.pos, direct):
            # print('1')
            self.qtable[x][y][direct] += self.alpha * \
                (self.reward[x][y] + self.gamma *
                 self.border - self.qtable[x][y][direct])
            return
        else:  # 其他合法情况
            if direct == 0:
                nx, ny = x-1, y
            elif direct == 1:
                nx, ny = x+1, y
            elif direct == 2:
                nx, ny = x, y-1
            elif direct == 3:
                nx, ny = x, y+1
            # 更新qtable
            self.qtable[x][y][direct] += self.alpha * \
                (self.reward[nx][ny] + self.gamma *
                 max(self.qtable[nx][ny]) - self.qtable[x][y][direct])
            npos = (nx, ny)
            # 如果下一个位置是可到达的，即不是墙壁或者别的agent
            if self.env.reachable(npos):
                # print('2')
                self.env.updatePos(self.pos, npos)
                self.pos = npos
            # print(self.qtable[x][y])

    def train(self, loops) -> None:
        for i in range(loops):
            while True:
                if self.isEnd():
                    print('name:', self.name, '次数：', i, '步数：', self.steps)
                    self.log.write('步数：' + str(self.steps) + '\n')
                    self.reset()
                    break
                self.action()
                # sleep(0.1)
        self.writeQtable(self.name + '_qtable.csv')
